import numpy as np
import cv2
import pyrealsense2 as rs
import torch
import torch.backends.cudnn as cudnn
from pathlib import Path
from ultralytics import YOLO
from DecisionMaker import RCDecisionMaker
import argparse,time,os
from ultralytics.utils.checks import check_requirements
from ultralytics.utils.torch_utils import strip_optimizer
from ultralytics.utils import set_logging
from ultralytics.utils.plotting import Annotator,colors,output_to_target,Colors
from ultralytics.utils.ops import scale_coords
from ultralytics.utils.torch_utils import time_sync
from ultralytics.utils.torch_utils import select_device
from ultralytics.data.loaders import LoadStreams

def attempt_load(weights: str = None, map_location: torch.device = None):
    """
    加载 .pt 权重文件并返回一个已经切换到 eval() 模式的模型。

    Args:
        weights (str): 权重文件路径
        map_location (torch.device): 设备映射，例如 torch.device('cuda') 或 'cpu'。

    Returns:
        torch.nn.Module: 加载好的模型，已设置为 eval() 模式。
    """
    # 获取当前脚本的绝对路径
    current_script_path = os.path.abspath(__file__)  # /home/robocon/Foursteps/youbaote/youbaote_depth2.py
    
    # 定位到项目根目录(Foursteps/)
    project_root = os.path.dirname(os.path.dirname(current_script_path))  # 两次dirname得到Foursteps/
    
    # 构造默认权重路径
    default_weights = os.path.join(
        project_root, 
        'yolov8', 
        'runs', 
        'train',
        'exp9',
        'weights',
        'best.pt'
    )
    
    # 优先级: 用户指定 > 默认路径
    weights = weights if weights else default_weights
    # 加载 checkpoint
    ckpt = torch.load(weights, map_location=map_location)
    # 取出模型（Ultralytics YOLOv8 的 .pt 一般把模型保存在 'model' 键下）
    model = ckpt['model'].float().eval()
    return model


def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    """
    调整图像大小并进行填充，同时满足步长倍数的约束。
    
    参数:
    img (numpy.ndarray): 输入的图像。
    new_shape (int or tuple): 目标图像的大小，如果是整数，则表示正方形的边长；如果是元组，则为 (高度, 宽度)。
    color (tuple): 填充边界时使用的颜色，格式为 (B, G, R)。
    auto (bool): 是否使用最小矩形填充方式。
    scaleFill (bool): 是否拉伸图像以填充整个目标大小。
    scaleup (bool): 是否允许图像放大。
    stride (int): 步长，用于确保填充后的图像尺寸是步长的倍数。
    
    返回:
    tuple: 包含调整后的图像、缩放比例和填充量的元组。
    """
    # 获取当前图像的高度和宽度
    shape = img.shape[:2]  # current shape [height, width]
    
    # 如果 new_shape 是整数，则将其转换为 (new_shape, new_shape) 的元组
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # 计算缩放比例（新尺寸 / 旧尺寸）
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    
    # 如果不允许图像放大，则取缩放比例和 1.0 中的较小值
    if not scaleup:  # only scale down, do not scale up (for better test mAP)
        r = min(r, 1.0)

    # 计算宽度和高度的缩放比例
    ratio = r, r  # width, height ratios
    
    # 计算不填充时的新宽度和高度
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    
    # 计算需要填充的宽度和高度
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    
    # 如果使用最小矩形填充方式
    if auto:  # minimum rectangle
        # 确保填充量是步长的倍数
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    
    # 如果使用拉伸填充方式
    elif scaleFill:  # stretch
        # 不进行填充
        dw, dh = 0.0, 0.0
        # 新尺寸即为目标尺寸
        new_unpad = (new_shape[1], new_shape[0])
        # 重新计算宽度和高度的缩放比例
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    # 将填充量平均分配到图像的两侧
    dw /= 2  # divide padding into 2 sides
    dh /= 2

    # 如果图像尺寸需要调整，则进行调整
    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    
    # 计算上下左右的填充量
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    
    # 使用指定颜色填充图像边界
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    
    # 返回调整后的图像、缩放比例和填充量
    return img, ratio, (dw, dh)

def process_img(img,device='cuda'):
    """
    处理输入图像，将其转换为适合模型推理的格式。

    参数:
    img (numpy.ndarray): 输入的图像数组，通常为OpenCV读取的BGR格式。
    device (torch.device): 计算设备，默认为 'cuda'（GPU）。

    返回:
    torch.Tensor: 处理后的图像张量。
    """
    # 判断是否使用半精度浮点数（fp16），如果设备不是CPU，则使用半精度
    half = device.type != 'cpu'
    # 调用 letterbox 函数调整图像大小并填充，不过这里未接收返回值，存在逻辑问题
    img ,ratio , pad = letterbox(img)
    # 将图像从BGR格式转换为RGB格式，并调整维度顺序为 (C, H, W)
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    # 将图像数组转换为内存连续的数组，以提高后续操作的效率
    img = np.ascontiguousarray(img)
    # 将NumPy数组转换为PyTorch张量，并移动到指定设备上
    img = torch.from_numpy(img).to(device)
    # 根据是否使用半精度，将张量数据类型转换为fp16或fp32
    img = img.half() if half else img.float()  # uint8 to fp16/32
    # 将图像像素值从 0 - 255 归一化到 0.0 - 1.0
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    # 如果图像张量维度为3，添加一个批次维度，使其维度变为 (1, C, H, W)
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    
    return img,ratio,pad


def detect(save_img=False):  
    """
    执行目标检测任务，支持普通摄像头和 RealSense 摄像头。

    参数:
    save_img (bool): 是否保存检测结果图像，默认为 False。
    """
    # 初始化决策
    RCSmartestBrain = RCDecisionMaker(OpenSerial=True,bufferSize=10)
    frameCounter = 0
    # 创建 RealSense 管道，用于管理摄像头数据流
    pipeline = rs.pipeline() 
    config = rs.config()
    # 启用深度流，设置分辨率为 640x480，格式为 16 位深度图，帧率为 30fps
    config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
    # 启用彩色流，设置分辨率为 640x480，格式为 BGR8，帧率为 30fps
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

    # 启动摄像头数据流
    pipeline.start(config)
    print("[WONDERFUL INFO] camera ready.")
    # 指定将深度图像对齐到彩色图像
    align_to = rs.stream.color
    # 创建对齐对象，用于将深度图像与彩色图像对齐
    align = rs.align(align_to) 

    # 从命令行参数中获取模型权重、是否显示图像和图像尺寸
    weights,view_img,img_size= opt.weights,opt.view_img,opt.img_size
    # 设置日志记录
    set_logging()
    # 选择计算设备，优先使用 GPU
    assert torch.cuda.is_available(), "cuda unavailable!"
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    print(f'device : ',device)
    # 加载模型并指定计算设备
    # model = attempt_load(weights=None, map_location=device)
    model = YOLO(opt.weights).to(device)

    # 获取模型的最大步长
    stride = int(model.stride.max())
    # 如果使用 GPU，将模型转换为半精度浮点数以提高计算效率
    if device.type == 'cuda':
        model.half()

    # 获取模型的类别名称和对应的颜色列表
    names = model.module.names if hasattr(model, 'module') else model.names
    print("names:", names, type(names))  # 检查 names 类型
    # colors_list = [[int(x) for x in np.random.randint(0, 255,3)] for _ in names]
    colors_list = [tuple(map(int, np.random.randint(0, 256, 3))) for _ in range(len(names))]
    # 初始化视频路径和视频写入器
    video_path,video_writer = None,None
    # 判断是否使用普通摄像头
    webcam = opt.webcam
    if webcam:
        # 如果使用普通摄像头，设置显示图像为 True
        view_img = True
        # 开启 cuDNN 基准测试，加速固定尺寸图像的推理
        cudnn.benchmark = True  
        # 创建数据流加载器，用于加载普通摄像头的视频流
        dataset = LoadStreams(sources='0')
    else:
        # 如果不使用普通摄像头，设置保存图像为 True
        save_img = True

    # 记录检测开始时间
    t0 = time.time()
    try:
        while True:  
            frameCounter += 1
            if webcam:
                # 处理普通摄像头的视频流
                for image_path, img, im0s, vid_cap in dataset:
                    # 处理输入图像，转换为适合模型推理的格式
                    img_tensor,ratio,pad = process_img(img, device)
                    # 使用模型进行目标检测
                    results = model.predict(
                        source=img_tensor,
                        imgsz=opt.img_size,
                        augment=opt.augment,
                        device=device,
                        stream=False
                    )[0]  

                    # 用 Annotator 画框并叠加深度
                    annotator = Annotator(im0s, line_width=3,example=str(names))
                    if results.boxes:
                        for box in results.boxes:
                            xyxy = box.xyxy[0].cpu().numpy()
                            # 置信度
                            conf = float(box.conf[0].cpu().item())
                            # class i
                            cls_i = int(box.cls[0].cpu().item())
                            label = f"{names[cls_i]} {conf:.2f}"
                            annotator.box_label(xyxy, label, color=colors_list[cls_i])                   
                        
                        # 获取带有检测结果的图像
                        original_image = annotator.result()
                        # 显示处理后的图像
                        cv2.imshow(str(image_path), original_image)
                        # 按 'q' 键退出循环
                        if cv2.waitKey(1) & 0xFF == ord('q'):
                            raise StopIteration
            else:
                # 处理 RealSense 摄像头的视频流
                frames = pipeline.wait_for_frames()
                # 对获取的帧进行对齐处理
                aligned_frames = align.process(frames)
                # 获取对齐后的彩色帧
                color_frame = aligned_frames.get_color_frame()
                # 获取对齐后的深度帧
                depth_frame = aligned_frames.get_depth_frame()
                if not color_frame or not depth_frame:
                    # 如果彩色帧或深度帧无效，跳过本次循环
                    continue
                # 将深度帧数据转换为 NumPy 数组
                depth_image = np.asanyarray(depth_frame.get_data())
                # 将彩色帧数据转换为 NumPy 数组
                color_image = np.asanyarray(color_frame.get_data())
                # 复制原始彩色图像
                original_image = color_image.copy()
                # 处理输入图像，转换为适合模型推理的格式
                img_tensor,ratio,pad = process_img(color_image, device)
                # 使用模型进行目标检测
                results = model.predict(img_tensor, augment=opt.augment)

                # 定义 RealSense 摄像头图像的路径
                image_path = Path('realsense_frame.jpg')
                # 计算保存路径，如果 save_img 为 True 则保存到 output 目录
                save_path = os.path.join('output', image_path.name) if save_img else None
                # 初始化状态字符串
                status_str = ''
                annotator = Annotator(original_image,line_width=3,example=str(names))
 
                if results[0].boxes:
                    for box in results[0].boxes:
                        # xyxy, conf, cls 都在 boxes 对象里
                        xyxy = box.xyxy[0].cpu().numpy()          # [x1, y1, x2, y2]
                        conf = float(box.conf[0].cpu().item())
                        cls_i = int(box.cls[0].cpu().item())
                        label = f"{names[cls_i]} {conf:.2f}"
                        annotator.box_label(xyxy, label, color=colors_list[cls_i])
                        # 计算检测框的中心坐标
                        center_x = int((int(xyxy[0]) + int(xyxy[2])) / 2)
                        center_y = int((int(xyxy[1]) + int(xyxy[3])) / 2)
                        # 获取检测框中心的深度信息
                        zDepth = depth_frame.get_distance(center_x, center_y)
                        print(f'zDepth',zDepth)       
                        # 定义线条厚度和字体厚度
                        tl, tf = 3, max(2, 1)
                        # 在原始图像上绘制深度信息
                        cv2.putText(original_image, f"{round(zDepth * 100, 2)} cm",
                                    (center_x, center_y), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
                        # 决策中
                        RCSmartestBrain.MakeDecision(YOLOResultBox=box,names=names,frameCounter=frameCounter,zDepth=zDepth)
                else:
                    RCSmartestBrain.MakeDecision(YOLOResultBox=results[0].boxes,names=names,frameCounter=frameCounter,zDepth=0)
                original_image = annotator.result()
                # 显示处理后的图像
                cv2.imshow(str(image_path), original_image)
                # 按 'q' 键退出循环
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
    finally:
        if not webcam:
            print("[MYSTERIOUS INFO] camera stop.")
            # 停止 RealSense 摄像头数据流
            pipeline.stop()
        # 关闭所有 OpenCV 窗口
        cv2.destroyAllWindows()

if __name__ == '__main__':
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--weights', nargs='+',type=str, default='/home/robocon/Foursteps/yolov8/weights/best.pt', help='your yolov8n.pt path')
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
  parser.add_argument('--update', action='store_true', help='update all models')
  parser.add_argument('--view-img', action='store_true', help='display results')
  parser.add_argument('--save-img', action='store_true', help='save results to ... ')
  parser.add_argument('--load_img_data',action='store_true',default=False,help='Whether to load data from image.')
  parser.add_argument('--save-dir', action='store_true', help='img save directory')
  parser.add_argument('--webcam',action='store_true',default=False,help='use webcam instead of realsense.')
  parser.add_argument('--augment', action='store_true', default=False, help='image argumentation for prediction.')
  opt = parser.parse_args()
  print(opt)
  
  # automaticly find ultralytics/requirements.txt 检查当前环境中的包版本
  # check_requirements()  
  
  
  # 上下文管理器，用于禁用梯度计算，减少内存消耗，提高推理速度
  with torch.no_grad():
    # 检查命令行参数中是否指定了 --update 选项
    if opt.update:  # update all models (to fix SourceChangeWarning)
      # 遍历不同的模型权重文件
      for opt.weights in ['yolov8n.pt', 'yolov8s.pt', 'yolov8m.pt']:
        # 调用 strip_optimizer 函数去除模型文件中的优化器状态，解决源变更警告
        strip_optimizer(opt.weights)
    else:
      # 如果未指定 --update 选项，则调用 detect 函数进行目标检测
      detect()
